import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


img_dir = '../../../../large_data/CV2/lesson/Day07'
# img_name = 'hammer.jpg'
img_name = 'Moscow.jpeg'
# img_dir = '../../../../large_data/pic'
# img_name = 'DSC05039.JPG'
# img_name = 'DSC05022_1.JPG'
# img_name = 'dog.jpg'
# img_name = 'dog_bird.jpg'
# img_name = 'football.jpg'
# img_name = 'messi5.jpg'
img_path = os.path.join(img_dir, img_name)

spr = 3
spc = 4
spn = 0
plt.figure(figsize=[8, 6])

sep('load')
img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)
img_bgr = cv.imread(img_path, cv.IMREAD_COLOR)
# img = cv.resize(img, None, fx=0.1, fy=0.1, interpolation=cv.INTER_CUBIC)
print('original shape', img.shape)
H, W = img.shape
H2 = (H - 1) // 2
W2 = (W - 1) // 2
my_show_img(img, 'ori gray', cmap='gray')
my_show_img(img_bgr, 'original', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('hist by plot')
spn += 1
plt.subplot(spr, spc, spn)
for i, c in enumerate(list('bgr')):
    hist = cv.calcHist([img_bgr], [i], None, [256], [0, 256])
    plt.plot(hist, color=c, label=c)
    plt.legend()
    plt.title('hist by plot')

sep('hist by ploylines')
bg = np.zeros((256, 256, 3), dtype=np.uint8)
colors = [
    (255, 0, 0),
    (0, 255, 0),
    (0, 0, 255)
]
for i, c in enumerate(list('bgr')):
    hist = cv.calcHist([img_bgr], [i], None, [256], [0, 256]).ravel()
    cv.normalize(hist, hist, 0, 255, cv.NORM_MINMAX)
    hist = 255 - hist
    print_numpy_ndarray_info(hist, 'hist')
    pts = np.int32([(i, v) for i, v in enumerate(hist)])
    pts = np.expand_dims(pts, (0, 1))  # ATTENTION cv.ploylines needs expand_dims 0, 1 for vector of points
    print_numpy_ndarray_info(pts, 'pts')
    cv.polylines(bg, pts, False, colors[i], 2)
my_show_img(bg, 'hist by ploylines', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('equalize')
img2 = cv.equalizeHist(img)
img2 = img2.astype(np.uint8)
my_show_img(img2, 'equalize', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('2d hist')
hsv = cv.cvtColor(img_bgr, cv.COLOR_BGR2HSV)
hist2d = cv.calcHist([hsv], [0, 1], None, [180, 256], [0, 180, 0, 256])
cv.normalize(hist2d, hist2d, 0, 255, cv.NORM_MINMAX)
hist2d = hist2d.astype(np.uint8)
my_show_img(hist2d, 'hist2d', cmap='gray')

sep('roi')
roi = hsv[0:100, 511:W]
my_show_img(roi, 'roi', lambda x: cv.cvtColor(x, cv.COLOR_HSV2RGB))
hist_roi = cv.calcHist([roi], [0, 1], None, [180, 256], [0, 180, 0, 256])
B = cv.calcBackProject([hsv], [0, 1], hist_roi, [0, 180, 0, 256], 1)
print_numpy_ndarray_info(B, 'B')

sep('mask')
my_show_img(B, 'mask', cmap='gray')

sep('masked')
img_masked = cv.bitwise_and(img_bgr, img_bgr, mask=B)
my_show_img(img_masked, 'masked', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('mask filtered')
kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5))
print(kernel)
B = cv.filter2D(B, -1, kernel)
my_show_img(B, 'mask filtered', cmap='gray')

sep('masked')
img_masked = cv.bitwise_and(img_bgr, img_bgr, mask=B)
my_show_img(img_masked, 'masked', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
